import os
import numpy as np
import torch
import time

def make_oracle(model, xy, temperature=1.0, device='cuda:0'):
    num_nodes = len(xy)
    
    # Convert xy to a tensor and move it to the specified device
    xyt = torch.tensor(xy, dtype=torch.float32).to(device)[None]  # Add batch dimension
    
    with torch.no_grad():  # Inference only
        embeddings, _ = model.embedder(model._init_embed(xyt))
        # Compute keys, values for the glimpse and keys for the logits once as they can be reused in every step
        fixed = model._precompute(embeddings)
    
    def oracle(tour):
        with torch.no_grad():  # 不计算trainable参数的梯度
            # Input tour with 0-based indices; output vector with probabilities for locations not in tour
            tour = torch.tensor(tour, dtype=torch.long, device=device)
            
            if len(tour) == 0:
                step_context = model.W_placeholder
            else:
                step_context = torch.cat((embeddings[0, tour[0]], embeddings[0, tour[-1]]), -1)

            # Compute query = context node embedding, add batch and step dimensions (both 1)
            query = fixed.context_node_projected + model.project_step_context(step_context[None, None, :])

            # Create the mask and convert to bool if needed
            mask = torch.zeros(num_nodes, dtype=torch.bool, device=device)
            mask[tour] = 1
            mask = mask[None, None, :]  # Add batch and step dimension

            log_p, _ = model._one_to_many_logits(query, fixed.glimpse_key, fixed.glimpse_val, fixed.logit_key, mask)
            p = torch.softmax(log_p / temperature, -1)[0, 0]
            
        #     assert (p[tour] == 0).all()
        #     assert (p.sum() - 1).abs() < 1e-5  # Check sum close to 1
            
        return p.cpu().numpy()  # Move back to CPU and return as NumPy array

    return oracle  
    
def get_index(oracle,n):
    tour = []
    while(len(tour) < n):
        p = oracle(tour)
        i = np.argmax(p)
        tour.append(i)
    return tour
